#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
股票分析多智能体系统
数据收集员 → 数据分析师 → 投资决策官
"""

import os
from datetime import datetime
from pathlib import Path
from crewai import Crew, Process, Task

# 导入工具（会自动修复编码）
from utils import (
    ensure_utf8_output,
    STOCK_DATA_TOOLS,
    NEWS_TOOLS,
    ALL_TOOLS
)

# 导入智能体
from agents import (
    create_data_collector_agent,
    create_data_analyst_agent,
    create_decision_maker_agent
)

# 确保 UTF-8 输出
ensure_utf8_output()

# 注意：配置已写死在代码中，无需 .env 文件


def get_llm():
    """获取配置好的 LLM 实例"""
    default_config = {
        "model": "gpt-4o-mini",
        "api_key": "sk-123123123",
        "api_base": "https://api.123123.xyz/v1"
    }
    
    # 设置环境变量（CrewAI 需要）
    os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", default_config["api_key"])
    os.environ["OPENAI_API_BASE"] = os.getenv("OPENAI_API_BASE", default_config["api_base"])
    os.environ["OPENAI_BASE_URL"] = os.getenv("OPENAI_API_BASE", default_config["api_base"])
    
    # CrewAI 使用 litellm，需要 openai/ 前缀
    model_name = os.getenv("OPENAI_MODEL_NAME", default_config["model"])
    if not model_name.startswith("openai/"):
        model_name = f"openai/{model_name}"
    
    return model_name


def analyze_stock(stock_query: str) -> dict:
    """
    执行完整的股票分析流程
    
    Args:
        stock_query: 股票查询（如"贵州茅台"、"600519"）
    
    Returns:
        包含各阶段结果的字典
    """
    print("\n" + "=" * 80)
    print(f"📊 开始分析股票：{stock_query}")
    print("=" * 80)
    print()
    
    # 识别股票代码
    stock_code = stock_query
    if "茅台" in stock_query:
        stock_code = "600519"
    elif "浦发" in stock_query:
        stock_code = "600000"
    elif "平安银行" in stock_query:
        stock_code = "000001"
    
    print(f"识别股票代码: {stock_code}")
    print()
    
    # 获取 LLM
    llm = get_llm()
    
    print("=" * 80)
    print("S1: 智能体创建阶段 - 初始化多智能体系统")
    print("=" * 80)
    print()
    
    # 创建智能体（传入工具，使用 CrewAI 工具调用链）
    agents = {
        "collector": create_data_collector_agent(ALL_TOOLS, llm),
        "analyst": create_data_analyst_agent(llm),
        "decision_maker": create_decision_maker_agent(llm)
    }
    
    print("✅ 已创建智能体：")
    print("  🔍 数据收集员 - 配备了 6 个工具（股票数据+新闻搜索）")
    print("  📊 数据分析师 - 负责分析股票")
    print("  💼 投资决策官 - 负责做出决策")
    print()
    
    # 创建任务（让 Agent 自主调用工具获取数据）
    task_collect = Task(
        description=f"""
请收集股票 {stock_query} (代码: {stock_code}) 的完整数据。

**你必须使用以下工具来获取数据：**
1. 使用 "获取股票实时价格" 工具获取实时行情
2. 使用 "获取股票基本信息" 工具获取公司信息
3. 使用 "获取股票历史数据" 工具获取近5天的K线数据（参数: stock_code="{stock_code}", days=5）
4. 使用 "搜索股票新闻" 工具获取最新新闻（参数: keyword="{stock_query}", max_results=3）

**输出要求：**
1. 整理工具返回的所有数据
2. 使用清晰的 Markdown 格式
3. 包含：股票名称、当前价格、涨跌幅、行业、近期走势、最新动态
4. 控制在 300 字以内

**重要：** 
- 必须调用所有 4 个工具获取完整数据
- 不要编造任何数据，只使用工具返回的真实数据
        """,
        agent=agents["collector"],
        expected_output="300字以内的 Markdown 格式数据摘要，基于工具返回的真实数据"
    )
    
    task_analyze = Task(
        description="""
基于收集的数据进行简要分析。

**输出要求：**
- 使用 Markdown 格式
- 控制在 200 字以内
- 包含：技术面评价、基本面评价、风险提示
        """,
        agent=agents["analyst"],
        expected_output="200字以内的 Markdown 格式分析报告",
        context=[task_collect]
    )
    
    task_decision = Task(
        description="""
给出最终投资建议。

**输出要求：**
- 使用 Markdown 格式
- 控制在 200 字以内
- 包含：操作建议（买入/持有/卖出）、目标价位、止损价位、核心理由
        """,
        agent=agents["decision_maker"],
        expected_output="200字以内的 Markdown 格式投资决策",
        context=[task_collect, task_analyze]
    )
    
    # 创建 Crew
    crew = Crew(
        agents=list(agents.values()),
        tasks=[task_collect, task_analyze, task_decision],
        process=Process.sequential,
        verbose=True
    )
    
    print("=" * 80)
    print("S2: 智能体协作阶段 - 执行多智能体任务流")
    print("=" * 80)
    print()
    print("🚀 开始执行协作任务...")
    print("   数据收集员将自主调用工具获取数据...")
    print("-" * 80)
    print()
    
    # 执行任务
    result = crew.kickoff()
    
    # 提取各阶段结果
    stages = {
        "stock_query": stock_query,
        "stock_code": stock_code,
        "stage1_data": str(task_collect.output.raw) if task_collect.output else "无数据",
        "stage2_analysis": str(task_analyze.output.raw) if task_analyze.output else "无分析",
        "stage3_decision": str(task_decision.output.raw) if task_decision.output else "无决策",
        "final_output": str(result)
    }
    
    return stages


def save_result(stock_query: str, stages: dict):
    """
    保存结果到 Markdown 文件
    
    Args:
        stock_query: 股票查询
        stages: 各阶段结果
    """
    # 创建 result 目录
    result_dir = Path("result")
    result_dir.mkdir(exist_ok=True)
    
    # 生成文件名
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    filename = f"stock_analysis_{timestamp}.md"
    filepath = result_dir / filename
    
    # 生成 Markdown 内容
    md_content = f"""# {stock_query} - 股票分析报告

**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}  
**股票代码**: {stages.get('stock_code', 'N/A')}

---

## 📋 多智能体协作流程

本报告由三个智能体通过 **CrewAI 工具调用链** 协作完成：

1. **数据收集员** - 使用 6 个工具自主收集股票实时数据、历史数据和新闻
2. **数据分析师** - 基于收集的数据分析技术面和基本面
3. **投资决策官** - 综合分析给出投资建议

---

## 阶段 1：数据收集（工具调用）

数据收集员通过 CrewAI 工具调用链自主获取数据：

{stages.get('stage1_data', '无数据')}

---

## 阶段 2：数据分析

{stages.get('stage2_analysis', '无分析')}

---

## 阶段 3：投资决策

{stages.get('stage3_decision', '无决策')}

---

## 💡 技术说明

- **工具调用方式**: CrewAI Tools (使用 @tool 装饰器)
- **工具列表**: 
  - 获取股票实时价格
  - 获取股票基本信息
  - 获取股票历史数据
  - 搜索股票新闻
  - 搜索市场新闻
  - 搜索行业新闻
- **调用方式**: Agent 根据任务需求自主选择和调用工具
- **数据来源**: Tushare (股票数据) + Tavily (新闻搜索)

---

*本报告由 CrewAI 多智能体系统 + 工具调用链自动生成*
"""
    
    # 保存文件
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(md_content)
    
    print(f"\n💾 完整结果已保存到: {filepath}")
    
    # 同时保存到 output.md（兼容旧版）
    output_file = Path("output.md")
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(md_content)
    
    print(f"💾 结果也已保存到: {output_file}")


def main():
    """主函数"""
    print("=" * 80)
    print("📈 股票分析多智能体系统")
    print("=" * 80)
    print()
    print("本系统演示四个阶段的协作流程：")
    print("  S1: 数据收集 → S2: 智能体创建 → S3: 智能体协作 → S4: 结果保存")
    print()
    
    # 定义要分析的股票
    stock_query = "贵州茅台"  # 可以改为 "600519" 或其他股票
    
    try:
        # 执行分析
        stages = analyze_stock(stock_query)
        
        print("\n" + "=" * 80)
        print("S3: 结果保存阶段 - 生成分析报告")
        print("=" * 80)
        print()
        
        # 保存结果
        save_result(stock_query, stages)
        
        print("\n" + "=" * 80)
        print("✅ 所有任务执行完成！")
        print("=" * 80)
        print("\n💡 技术亮点：Agent 通过 CrewAI 工具调用链自主获取数据")
        print("   不再是手动调用工具后将数据注入，而是让 Agent 自己决定何时调用哪个工具")
        
        # 显示最终决策结果
        print("\n" + "=" * 80)
        print("S4: 最终输出阶段 - 投资决策建议")
        print("=" * 80)
        print()
        print(stages.get('stage3_decision', '无决策'))
        print()
        print("=" * 80)
        
    except Exception as e:
        print(f"\n❌ 执行出错：{str(e)}")
        print("请检查 API 配置是否正确")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
